from contextlib import asynccontextmanager
from functools import wraps

from tornado.websocket import WebSocketHandler as OriginWebsocketHandler
from tornado.websocket import WebSocketClosedError
from tornado.ioloop import IOLoop
from marshmallow import ValidationError

from webutils.compat import complex_json
from webutils.exceptions import UserRaiseAbort
from webutils.layers.connection import ConnectionProxy
from webutils.logging import get_summary, LoggingMixin
from webutils.ioloop import BackgroundTaskMixin
from webutils.datastructure import NameSpace

from .mixin import RemoteIPMixin
from .responses import Response


def callback_wrapper(f):
    @wraps(f)
    async def wrapper(handler, *args, **kwargs):
        try:
            return await f(handler, *args, **kwargs)
        except Exception as exc:
            return await handler.handle_callback_exception(exc)

    return wrapper


class WebSocketHandler(BackgroundTaskMixin, LoggingMixin, RemoteIPMixin,
                       OriginWebsocketHandler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.background_tasks = {}
        self.connection_session = NameSpace(_allow_update=True)

    @property
    def session(self):
        return self.connection_session

    @property
    def is_live(self):
        return self.ws_connection and not self.ws_connection.is_closing()

    def log_exception(self, typ, exc, tb):
        self.stop_background_tasks()
        IOLoop.current().add_callback(self.release_resource)
        return super().log_exception(typ, exc, tb)

    async def handle_callback_exception(self, exc):
        if isinstance(exc, UserRaiseAbort):
            return self.close()
        elif isinstance(exc, ValidationError):
            msg = str(exc)
            resp = Response(errcode='500', errmsg='invalid parameters: ' + msg)
            await self.write_response(resp)
            return
        raise exc

    async def write_message(self, *args, **kwargs):
        if not self.is_live:
            self.logger.info('try to write_message while websocket is closed')
            return
        try:
            await super().write_message(*args, **kwargs)
        except WebSocketClosedError:
            self.logger.warning(
                'failed to write message because websocket is closed')

    async def write_json_message(self, content):
        dumped_message = complex_json.dumps(content)
        return await self.write_message(dumped_message)

    async def write_response(self, response):
        assert isinstance(response, Response)
        return await self.write_json_message(response)

    async def release_resource(self):
        pass

    def on_close(self):
        IOLoop().current().add_callback(self.release_resource)

    def abort(self, msg):
        self.logger.info('abort because of %s', msg)
        raise UserRaiseAbort

    @asynccontextmanager
    async def get_connection(self):
        async with ConnectionProxy(self.application,
                                   log_context=self.log_context) as conn:
            yield conn

    @asynccontextmanager
    async def service(self, service_cls, *args, **kwargs):
        async with self.get_connection() as conn:
            kwargs.setdefault('log_context', self.log_context)
            yield service_cls(conn, *args, **kwargs)

    def _request_summary(self):
        origin_summary = super()._request_summary()
        extra_summary = get_summary(self.log_context)
        return f'{origin_summary} {extra_summary}'
